import yaml
import subprocess
import os

def main():
    # Load config
    config_path = 'scripts/eval.yaml'
    with open(config_path, 'r') as file:
        config = yaml.safe_load(file)

    # Evaluate
    weight_path = config['test']['weight_path']
    batch_size = config['test']['batch_size']
    gpu_num = config['test']['gpu_num']
    data_name = config['test']['data_name']
    exp_name = config['test']['exp_name']
    model = config['test']['model']
    visual_sample_count = config['test']['visual_sample_count']
    eval_name = config['test']['eval_name']
    
    if model == 'MatrixVT':
        work_dir = '.'
        command = [
            'python', 'bevdepth/exps/nuscenes/MatrixVT/matrixvt_bev_depth_lss_r50_256x704_128x128_24e_ema.py',
            '--exp', exp_name,
            '--data_root', data_name,
            '--ckpt_path', weight_path,
            '-e',
            '-b', str(batch_size),
            '--gpus', str(gpu_num)
        ]
        ret = subprocess.run(command, cwd=work_dir)
        if ret.returncode == 0:
            print("Finish Inference!")
        else:
            print("Error:", ret)
    else:
        raise ValueError(f'Model {model} not supported yet!') 
    
    # Convert Labels
    print('Start Evaluation')
    gt_path = os.path.join('../data', data_name, 'nuscenes/nuscenes_infos_val.pkl')
    pred_path = os.path.join('../outputs/', exp_name, 'results_nusc.json')
    work_dir = 'bev_eval_visual'
    command = [
        'python', 'tools/label_converter.py',
        '--gt_path', gt_path,
        '--pred_path', pred_path,
        '--eval_name', eval_name  
    ]
    ret = subprocess.run(command, cwd=work_dir)
    if ret.returncode == 0:
        print("Finish Label Convert!")
    else:
        print("Error:", ret)
    
    # Calculate metrics
    command = [
        'python', 'tools/nuscenes_eval.py',
        '--pred_labels', os.path.join(eval_name, 'pred_labels'),
        '--gt_labels', os.path.join(eval_name, 'gt_labels'),
        '--eval_name', eval_name,
        '--format',  'class x y z l w h r score' 
    ]
    ret = subprocess.run(command, cwd=work_dir)
    if ret.returncode == 0:
        print("Finish Evaluation!")
    else:
        print("Error:", ret)
    
    # # Visualization
    print('Start Visualization')
    command = [
        'python', 'tools/visualizer.py',
        '--data_name', data_name,
        '--exp_name', exp_name,
        '--eval_name', eval_name,
        '--sample_num', str(visual_sample_count)
    ] 
    ret = subprocess.run(command, cwd=work_dir)
    if ret.returncode == 0:
        print("Finish Visualization!")
    else:
        print("Error:", ret)

if __name__ == '__main__':
    main()
